/*
 * maximum_likelihood_trainer.h
 *
 *  Created on: Feb 3, 2012
 *      Author: kkb110
 */

#ifndef MAXIMUM_LIKELIHOOD_TRAINER_H_
#define MAXIMUM_LIKELIHOOD_TRAINER_H_

#include "common_includes.h"

class MaximumLikelihoodTrainer{
private:
	Model<double, 2, 2>::Shot mean(vector<Model<double, 2, 2>::Shot> shots){
		Model<double, 2, 2>::Shot result;
		result << 0, 0;
		for (size_t i = 0; i < shots.size(); ++i){
			result += shots[i];
		}
		return result / shots.size();
	}
	Eigen::Matrix<double, 2, 2> cov(vector<Model<double, 2, 2>::Shot> shots){
		Model<double, 2, 2>::Shot mean = this->mean(shots);
		double cov11 = 0;
		double cov12 = 0;
		double cov22 = 0;
		for (size_t i = 0; i < shots.size(); ++i){
			auto tmp = shots[i] - mean;
			cov11 += tmp(0) * tmp(0);
			cov12 += tmp(0) * tmp(1);
			cov22 += tmp(1) * tmp(1);
		}
		cov11 /= shots.size();
		cov12 /= shots.size();
		cov22 /= shots.size();
		Eigen::Matrix<double, 2, 2> result;
		result << cov11, cov12, cov12, cov22;
		return result;
	}
public:
	vector<vector<Shooter>> shooters;
	vector<vector<Model<double, 2, 2>::Shot>> shots;

	Model<double, 2, 2> run(){
		vector<Model<double, 2, 2>::Shot> shots1;
		vector<Model<double, 2, 2>::Shot> shots2;
		for (size_t k = 0; k < shooters.size(); ++k){
			for (size_t i = 0; i < shooters[k].size(); ++i){
				switch (shooters[k][i]) {
				case 0:
					shots1.push_back(shots[k][i]);
					break;
				case 1:
					shots2.push_back(shots[k][i]);
					break;
				}
			}
		}

		Model<double, 2, 2> current_guess;

		//count transition number
		current_guess.transition_probabilities << 0, 0, 0, 0;
		for (size_t k = 0; k < shooters.size(); ++k){
			for (size_t i(0); i < shooters[k].size() - 1; ++i){
				current_guess.transition_probabilities(shooters[k][i],
						shooters[k][i + 1]) += 1;
			}
		}
		for (size_t i(0); i < current_guess.transition_probabilities.rows();
				++i){
			double sum = 0.0;
			for (size_t j(0); j < current_guess.transition_probabilities.cols();
					++j){
				sum += current_guess.transition_probabilities(i, j);
			}
			for (size_t j(0); j < current_guess.transition_probabilities.cols();
					++j){
				current_guess.transition_probabilities(i, j) /= sum;
			}
		}

		//gaussin fitting (EM?)

		current_guess.shooting_distributions[0].mean = mean(shots1);
		current_guess.shooting_distributions[0].covariance = cov(shots1);
		current_guess.shooting_distributions[1].mean = mean(shots2);
		current_guess.shooting_distributions[1].covariance = cov(shots2);

//		cout << current_guess.shooting_distributions[0].mean << endl;
//		cout << current_guess.shooting_distributions[0].covariance << endl;
//		cout << current_guess.shooting_distributions[1].mean << endl;
//		cout << current_guess.shooting_distributions[1].covariance << endl;

		return current_guess;
	}

};

#endif /* MAXIMUM_LIKELIHOOD_TRAINER_H_ */
